package cn.itcast.manus.agent;

import cn.itcast.manus.agent.amap.AMAPAgent;
import cn.itcast.manus.agent.browser.ReActBrowserAgent;
import cn.itcast.manus.agent.chart.ChartAgent;
import cn.itcast.manus.agent.db.DBAgent;
import cn.itcast.manus.agent.html.HtmlDocAgent;
import cn.itcast.manus.agent.markdown.MarkdownAgent;
import cn.itcast.manus.agent.planning.ReActPlanningAgent;
import cn.itcast.manus.agent.table.TableAgent;
import cn.itcast.manus.enums.AgentTypeEnum;
import cn.itcast.manus.message.MessageSession;
import jakarta.annotation.PostConstruct;
import org.springframework.beans.factory.config.ConfigurableBeanFactory;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Scope;

import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;

/*
 * 当前实现中，AgentFactory 只注册了 RE_ACT_PLANNING_AGENT 类型的智能体，
 * 但系统设计上支持扩展到 AgentTypeEnum 中定义的所有智能体类型（如浏览器智能体、表格智能体等）。
 * 只需在 init() 方法中添加对应的注册代码即可快速扩展系统功能。
 */
@Configuration
public class AgentFactory {

    public static final Map<AgentTypeEnum, Function<MessageSession, Agent>> AGENT_FUNC_MAP = new HashMap<>();

    /**
     * 初始化方法，完成Agent的注册
     */
    @PostConstruct
    public void init() {
        // 注册任务规划智能体
        AGENT_FUNC_MAP.put(AgentTypeEnum.RE_ACT_PLANNING_AGENT, this::reActPlanningAgent);
        // 注册浏览器智能体
        AGENT_FUNC_MAP.put(AgentTypeEnum.BROWSER_AGENT, this::reActBrowserAgent);
        // 注册表格智能体
        AGENT_FUNC_MAP.put(AgentTypeEnum.TABLE_AGENT, this::tableAgent);
        // 注册Markdown智能体
        AGENT_FUNC_MAP.put(AgentTypeEnum.MARKDOWN_AGENT, this::markdownAgent);
        // 注册图表智能体
        AGENT_FUNC_MAP.put(AgentTypeEnum.CHART_AGENT, this::chartAgent);
        // 注册网页文档智能体
        AGENT_FUNC_MAP.put(AgentTypeEnum.HTML_DOC_AGENT, this::htmlDocAgent);
        // 注册高德地图智能体
        AGENT_FUNC_MAP.put(AgentTypeEnum.AMAP_AGENT, this::amapAgent);
        // 注册Mysql智能体
        AGENT_FUNC_MAP.put(AgentTypeEnum.DB_AGENT, this::dBAgent);

    }

    @Bean
    @Scope(value = ConfigurableBeanFactory.SCOPE_PROTOTYPE)
    public Agent dBAgent(MessageSession messageSession) {
        return new DBAgent(messageSession);
    }


    @Bean
    @Scope(value = ConfigurableBeanFactory.SCOPE_PROTOTYPE)
    public Agent amapAgent(MessageSession messageSession) {
        return new AMAPAgent();
    }

    @Bean
    @Scope(value = ConfigurableBeanFactory.SCOPE_PROTOTYPE)
    public Agent htmlDocAgent(MessageSession messageSession) {
        return new HtmlDocAgent(messageSession);
    }

    @Bean
    @Scope(value = ConfigurableBeanFactory.SCOPE_PROTOTYPE)
    public Agent chartAgent(MessageSession messageSession) {
        return new ChartAgent(messageSession);
    }

    @Bean
    @Scope(value = ConfigurableBeanFactory.SCOPE_PROTOTYPE)
    public Agent markdownAgent(MessageSession messageSession) {
        return new MarkdownAgent(messageSession);
    }

    @Bean
    @Scope(value = ConfigurableBeanFactory.SCOPE_PROTOTYPE)
    public Agent tableAgent(MessageSession messageSession) {
        return new TableAgent(messageSession);
    }

    @Bean
    @Scope(value = ConfigurableBeanFactory.SCOPE_PROTOTYPE)
    public Agent reActBrowserAgent(MessageSession messageSession) {
        return new ReActBrowserAgent(messageSession);
    }
    /**
     * 根据agentTypeEnum获取对应的Agent
     */
    public static Function<MessageSession, Agent> getAgent(AgentTypeEnum agentTypeEnum) {
        Function<MessageSession, Agent> fun = AGENT_FUNC_MAP.get(agentTypeEnum);
        if (null == fun) {
            throw new IllegalArgumentException("找不到对应的智能体: " + agentTypeEnum);
        }
        return fun;
    }

    /**
     * 任务规划智能体，交由Spring管理
     * ConfigurableBeanFactory.SCOPE_PROTOTYPE注解确保每次请求都会创建一个新的智能体实例，这对于多线程环境下的会话隔离非常重要：
     * @param messageSession 会话对象
     * @return 查找到的智能体实例
     */
    @Bean
    @Scope(value = ConfigurableBeanFactory.SCOPE_PROTOTYPE)
    public Agent reActPlanningAgent(MessageSession messageSession) {
        return new ReActPlanningAgent(messageSession);
    }

}